#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import re


def getOrganization():
    import pymysql

    # MySQL Info
    DATABASE_NAME = 'plosp'
    HOST = 'localhost'
    PORT = 3306
    USER_NAME = 'plosp'
    PASSWORD = 'plosp'
    CHAR_SET = 'utf8mb4'

    conn = pymysql.connect(host=HOST, port=PORT, user=USER_NAME, password=PASSWORD, db=DATABASE_NAME, charset=CHAR_SET)
    cursor = conn.cursor()
    sql = "SELECT id,name,description FROM tb_dds_organization"
    try:
        cursor.execute(sql)
        for row in cursor.fetchall():
            id = str(row[0])
            name = row[1].replace("\n", "")
            des = row[2].replace("\n", "")

            print(id + "\t" + name + "\t" + des)
    except:
        print("Error: unable to fetch data.")
    cursor.close()
    conn.close()


def SplitOrganization(filepath):
    count = 0
    infile = open(filepath, 'r')
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        if len(splitline) >= 2:
            id = splitline[0]
            des = splitline[2]
            regex = '<aff.*?>(.*?)</aff>'
            regex2 = '<label>.*?</label>'
            pattern = re.compile(regex)
            match = pattern.findall(des)
            for item in match:
                # count_label=item.count("<label>")
                column = re.split(regex2, item)
                for item2 in column:
                    item2 = item2.strip()
                    if item2.startswith("and"):
                        item2 = item2[3:].strip()
                    elif item2.endswith("and"):
                        item2 = item2[:-3].strip()
                    if item2.strip() != "":
                        count += 1
                        print(str(count) + "\t" + id + "\t" + item2)
    infile.close()

def CleanOrganization(org_string):
    from flashtext.keyword import KeywordProcessor
    keyword_processor = KeywordProcessor()
    keyword_processor.add_keyword('Big Apple', 'New York')
    keyword_processor.add_keyword('New Delhi', 'NCR region')
    new_sentence = keyword_processor.replace_keywords('I love Big Apple and new delhi.')


def GetExactName(name):
    regex = "\W(.*)\W"
    pattern = re.compile(regex, re.I)
    m = re.search(pattern, name)
    if m:
        return m.group(1)
    else:
        return m


def Search_CaseSensitive(regex, line):
    pattern = re.compile(regex)
    m = re.search(pattern, line)
    if m:
        return True
    else:
        return False


def Search_CaseIgnore(regex, line):
    pattern = re.compile(regex, re.I)
    m = re.search(pattern, line)
    if m:
        return True
    else:
        return False


def Search_CaseSensitive_Return(regex, line):
    regex2 = "(" + regex + ")"
    pattern = re.compile(regex2)
    m = re.search(pattern, line)
    if m:
        return m.group(1)
    else:
        return m


def Search_CaseIgnore_Return(regex, line):
    regex2 = "(" + regex + ")"
    pattern = re.compile(regex2, re.I)
    m = re.search(pattern, line)
    if m:
        return m.group(1)
    else:
        return m


def Initialize_Country_Regex(filepath):
    infile = open(filepath, "r")
    country_regex = ""
    country_abbr_regex = ""
    country_ht = {}
    country_abbr_ht = {}
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        country_regex += "\W" + splitline[1] + "\W|"
        country_ht[splitline[1].upper()] = splitline[1]
        if len(splitline) > 3:
            splitcountry = splitline[3].split("|")
            for item in splitcountry:
                country_abbr_regex += "\W" + item + "\W|"
                country_abbr_ht[item] = splitline[1]
    infile.close()
    country_regex = country_regex[:-1].replace(".", "\.")
    country_abbr_regex = country_abbr_regex[:-1].replace(".", "\.")
    return country_regex, country_abbr_regex, country_ht, country_abbr_ht


def Initialize_State_Regex(filepath):
    infile = open(filepath, "r")
    state_regex = ""
    state_abbr_regex = ""
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        state_regex += "\W" + splitline[1] + "\W|\W" + splitline[3] + "\W|"
        splitstate = splitline[2].split("|")
        for item in splitstate:
            state_abbr_regex += "\W" + item + "\W|"
    infile.close()
    state_regex = state_regex[:-1].replace(".", "\.")
    state_abbr_regex = state_abbr_regex[:-1].replace(".", "\.")
    return state_regex, state_abbr_regex


def Initialize_Other_Regex(filepath):
    infile = open(filepath, "r")
    other_regex = ""
    other_abbr_regex = ""
    other_ht = {}
    other_abbr_ht = {}
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        other_regex += "\W" + splitline[1] + "\W|"
        other_ht[splitline[1].upper()] = splitline[0]
        if len(splitline) > 2:
            splitother = splitline[2].split("|")
            for item in splitother:
                other_abbr_regex += "\W" + item + "\W|"
                other_abbr_ht[item] = splitline[0]
    infile.close()
    other_regex = other_regex[:-1].replace(".", "\.")
    other_abbr_regex = other_abbr_regex[:-1].replace(".", "\.")
    return other_regex, other_abbr_regex, other_ht, other_abbr_ht


def FindLabel(filepath):
    import re

    ht = {}
    infile = open(filepath, 'r')
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        des = splitline[2]
        regex = '<label>(.*?)</label>'
        pattern = re.compile(regex)
        match = pattern.findall(des)
        for item in match:
            item = item.strip()
            if item != "":
                if ht.has_key(item):
                    ht[item] += 1
                else:
                    ht[item] = 1
    infile.close()
    for key in ht.keys():
        print(key + "====" + str(ht[key]))


def FilterXMLLabel(line):
    import re
    regex = '<.*?>'
    items = re.split(regex, line)
    result = ""
    for item in items:
        result += item
    return result


def FilterEmail(line):
    import re
    email_regex = '(.*?); E-?Mails?:'
    pattern = re.compile(email_regex, re.I)
    match = pattern.findall(line)
    if len(match) == 0:
        return line
    else:
        possible_country = ""
        for item in match:
            possible_country = item.strip()
            break
        return possible_country


def FilterZipcode(line):
    import re
    regex = '\s*[\d\-]{4,}\s*;?'
    items = re.split(regex, line)
    result = ""
    for item in items:
        result += item
    return result


def StatCountry(filepath):
    ht = {}
    infile = open(filepath, 'r')
    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        des = splitline[2]
        des = FilterXMLLabel(des)
        if des.endswith(",") or des.endswith("."):
            des = des[:-1].rstrip()
        if des.count(",") > 0:
            possible_country = des.split(",")[-1].strip()
            possible_country = FilterEmail(possible_country)
            possible_country = FilterZipcode(possible_country)

            if ht.has_key(possible_country):
                ht[possible_country] += 1
            else:
                ht[possible_country] = 1
    infile.close()
    for key in ht.keys():
        print(key + "\t[" + str(ht[key]) + "]")


def RecognizeCountryStateCity(filepath, regexpath):
    infile = open(filepath, "r")
    outfile = open(filepath + ".miss", "w")
    outfile2 = open(filepath + ".onlyorg", "w")

    country_regex, country_abbr_regex = Initialize_Country_Regex(regexpath + "/country.compact.tsv")
    state_regex, state_abbr_regex = Initialize_State_Regex(regexpath + "/state.tsv")
    other_regex = Initialize_Other_Regex(regexpath + "/location.tsv")
    # print country_regex
    # print country_abbr_regex
    # print state_regex
    # print state_abbr_regex
    # print other_regex

    total_count = 0
    country_count = 0
    state_count = 0
    other_count = 0

    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        id = splitline[0]
        org = splitline[2]
        if Search_CaseIgnore(country_regex, org):
            total_count += 1
            country_count += 1
        elif Search_CaseSensitive(country_abbr_regex, org):
            total_count += 1
            country_count += 1
        elif Search_CaseIgnore(state_regex, org):
            total_count += 1
            state_count += 1
        elif Search_CaseSensitive(state_abbr_regex, org):
            total_count += 1
            state_count += 1
        elif Search_CaseIgnore(other_regex, org):
            total_count += 1
            other_count += 1
        else:
            org = org.rstrip(", ")
            if org.find(",") == -1:
                outfile2.write(line + "\n")
                outfile2.flush()
            else:
                outfile.write(line + "\n")
                outfile.flush()
    infile.close()
    outfile.close()
    outfile2.close()
    print("TOTAL: " + str(total_count))
    print("COUNTRY: " + str(country_count))
    print("STATE: " + str(state_count))
    print("LOCATION: " + str(other_count))


def GetCity(filepath):
    ht = {}
    infile = open(filepath, "r")
    for line in infile:
        line = line.rstrip()
        org = line.split("\t")[2]
        org = org.rstrip(";., ")
        splitcomma = org.split(",")
        possiblecity = splitcomma[-1].strip()
        splitcity = possiblecity.split(" ")
        if len(splitcity) == 1:
            if ht.has_key(possiblecity):
                ht[possiblecity] += 1
            else:
                ht[possiblecity] = 1
    infile.close()
    for key in ht.keys():
        print(key + "\t" + str(ht[key]))


def RecognizeCountryStateCity2(filepath, regexpath):
    from pymongo import MongoClient

    DB_NAME = "organization"
    COLLECTION_NAME = "organization"

    connection = MongoClient('localhost', 27017)
    db = connection[DB_NAME]
    collection = db[COLLECTION_NAME]

    infile = open(filepath, "r")

    country_regex, country_abbr_regex, country_ht, country_abbr_ht = Initialize_Country_Regex(
        regexpath + "/country.compact.tsv")
    state_regex, state_abbr_regex = Initialize_State_Regex(regexpath + "/state.tsv")
    other_regex, other_abbr_regex, other_ht, other_abbr_ht = Initialize_Other_Regex(regexpath + "/location.tsv")

    for line in infile:
        line = line.rstrip()
        splitline = line.split("\t")
        id = int(splitline[0])
        org_id = int(splitline[1])
        org = splitline[2]

        current_record = {}
        current_record["id"] = id
        current_record["org_id"] = org_id
        current_record["org"] = org
        org = " " + org + " "

        m1 = Search_CaseIgnore_Return(country_regex, org)
        if m1 != None:
            country = country_ht[GetExactName(m1).upper()]
            current_record["country"] = country
        else:
            m2 = Search_CaseSensitive_Return(country_abbr_regex, org)
            if m2 != None:
                country = country_abbr_ht[GetExactName(m2)]
                current_record["country"] = country
            elif Search_CaseIgnore(state_regex, org):
                current_record["country"] = "United States"
            elif Search_CaseSensitive(state_abbr_regex, org):
                current_record["country"] = "United States"
            else:
                m3 = Search_CaseIgnore_Return(other_regex, org)
                if m3 != None:
                    country = other_ht[GetExactName(m3).upper()]
                    current_record["country"] = country
                else:
                    m4 = Search_CaseSensitive_Return(other_abbr_regex, org)
                    if m4 != None:
                        country = other_abbr_ht[GetExactName(m4)]
                        current_record["country"] = country
        collection.insert_one(current_record)
    infile.close()
    connection.close()


if __name__ == "__main__":
    import sys

    if sys.argv[1] == "getOrganization":
        getOrganization()
    elif sys.argv[1] == "SplitOrganization":
        SplitOrganization(sys.argv[2])
    elif sys.argv[1] == "FindLabel":
        FindLabel(sys.argv[2])
    elif sys.argv[1] == "StatCountry":
        StatCountry(sys.argv[2])
    elif sys.argv[1] == "Recognize":
        RecognizeCountryStateCity(sys.argv[2], sys.argv[3])
    elif sys.argv[1] == "GetCity":
        GetCity(sys.argv[2])
    elif sys.argv[1] == "Recognize2":
        RecognizeCountryStateCity2(sys.argv[2], sys.argv[3])
